import time

import cv2
import numpy as np
import matplotlib.pyplot as plt

class Window_3d:
    def __init__(self):
        super().__init__()
        self.fig = plt.figure(figsize=(6.4, 4.8))

        # 创建两个字图
        self.ax1 = self.fig.add_subplot(1, 2, 1, projection='3d')

        self.ax2 = self.fig.add_subplot(1, 2, 2, projection='3d')


    def update_win(self,pts=None,lines=None,length=200):
            pts = np.array(pts)
            if pts is not None and len(pts) > 0:
                step = 1
                if len(pts)>length:
                    step = len(pts)//length
                pts = pts[::step]
                print(len(pts))
                x, y, z = pts[:, 0], pts[:, 1], pts[:, 2]
                self.ax1.scatter(x, y, z)
            if lines is not None:
                lines = np.array(lines)
                # l = lines[0]
                for l in lines:
                    if len(l) > 0:
                        step = 1
                        if len(l) > length:
                            step = len(l) // length
                        l = l[::step]
                        # print(len(l))
                        # print(l[:,:-1, 3])
                        x, y, z = l[:,0],l[:,1],l[:,2]
                        xyz = np.column_stack((x,y,z))
                        # print(x,y,z)
                        self.ax2.plot(xyz[:,0],xyz[:,1],xyz[:,2], c = tuple(np.random.randint(0, 255, 3)/255))

                plt.show(block=False)
                plt.pause(0.01)
                plt.cla()

def draw_(limg,lpts,rimg,rpts):
    if len(lpts) == len(rpts):
        for (p1,p2) in zip(lpts,rpts):
            p1,p2 = p1.pt,p2.pt
            p1 = (int(p1[0]),int(p1[1]))
            p2 = (int(p2[0]),int(p2[1]))
            c = tuple(np.random.randint(0,255,3).tolist())
            cv2.circle(limg,p1,2,c,-1)
            cv2.circle(rimg,p2,2,c,-1)
    return limg,rimg
def read_imgs(win1,steroslam):
    ltxt_path = 'slam_dataset/binocular/TUM/skate-easy-vi_gt_data/left.txt'
    rtxt_path = 'slam_dataset/binocular/TUM/skate-easy-vi_gt_data/right.txt'
    lf = open(ltxt_path,'r')
    rf = open(rtxt_path,'r')

    l_list = lf.read().split('\n')
    l_list = sorted(l_list,key=lambda x:int(x.split('.')[0].split('/')[-1]))
    r_list = rf.read().split('\n')
    r_list = sorted(r_list,key=lambda x:int(x.split('.')[0].split('/')[-1]))

    path_tup = zip(l_list,r_list)
    for (pt1,pt2) in path_tup:
        pt1 = 'slam_dataset/'+'/'.join(pt1.split('/')[4:])
        pt2 = 'slam_dataset/'+'/'.join(pt2.split('/')[4:])
        # print(pt1,pt2)
        l_img = cv2.imread(pt1)
        r_img = cv2.imread(pt2)

        steroslam.update((l_img,r_img))
        steroslam.compute_xyz_camera()
        path1 = steroslam.lpath
        path2 = steroslam.rpath
        path1 = np.array(path1)
        path2 = np.array(path2)
        # print(path1,path2)
        # print(steroslam.map_pts)
        # print(steroslam.points3D)
        win1.update_win(steroslam.map_pts,[path1,path2])
        l_img,r_img = draw_(l_img,steroslam.prev_lkps,r_img,steroslam.prev_rkps)
        # print(len(pts))
        # print(lpts)
        img = np.hstack((l_img,r_img))
        # print(img.shape,img.dtype)
        cv2.imshow('img',img[::2,::2])
        cv2.waitKey(50)

class stero_SLAM:
    def __init__(self):
        super().__init__()
        self.prev_limg = None
        self.prev_rimg = None

        self.prev_ldes = None
        self.prev_rdes = None

        self.prev_lkps = None
        self.prev_rkps = None

        self.lk = None
        self.rk = None

        self.lpose = np.eye(4)
        self.rpose = np.eye(4)

        self.lpath = []
        self.rpath = []

        # 三维重建
        self.points3D = []
        self.points2D_left = []
        self.points2D_right = []
        self.map_pts = []
        self.orb = cv2.ORB_create()
        self.bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        self.set_k()
    def set_k(self):
        self.lk = np.array([[748.3, 0, 490.6],
                           [0, 748.3, 506.3],
                           [0, 0, 1]])

        self.rk = np.array([[743.7, 0, 495.3],
                        [0, 743.4, 514.7],
                        [0, 0, 1]])

    def update(self,frames):
        left_img,right_img = frames
        temps = time.time()
        left_img = cv2.cvtColor(left_img,cv2.COLOR_BGR2GRAY)
        right_img = cv2.cvtColor(right_img,cv2.COLOR_BGR2GRAY)

        lkps,ldes = self.orb.detectAndCompute(left_img, None)
        rkps,rdes = self.orb.detectAndCompute(right_img, None)
        if self.prev_lkps is not None:
            lmatches = self.bf.match(self.prev_ldes,ldes)
            N = 100
            lmatches = lmatches[:N]
            src_pts = np.float32([self.prev_lkps[m.queryIdx].pt for m in lmatches]).reshape(-1, 1, 2)
            dst_pts = np.float32([lkps[m.trainIdx].pt for m in lmatches]).reshape(-1, 1, 2)
            E, mask = cv2.findEssentialMat(dst_pts, src_pts, self.lk, method=cv2.RANSAC, prob=0.999, threshold=1.0)
            src_pts_filtered = src_pts[mask.ravel() == 1]
            dst_pts_filtered = dst_pts[mask.ravel() == 1]
            src_pts_filtered = np.round(src_pts_filtered).astype(np.int16).squeeze()
            dst_pts_filtered = np.round(dst_pts_filtered).astype(np.int16).squeeze()

            _, R, t, mask = cv2.recoverPose(E, dst_pts_filtered, src_pts_filtered, self.lk)

            self.lpose[:3, :3] = np.dot(R, self.lpose[:3, :3])
            self.lpose[:3, 3] += self.lpose[:3, :3].dot(t.squeeze())
            self.lpath.append(self.lpose[:3, 3].copy())

            rmatches = self.bf.match(self.prev_rdes,rdes)
            rmatches = rmatches[:N]
            src_pts = np.float32([self.prev_rkps[m.queryIdx].pt for m in rmatches]).reshape(-1, 1, 2)
            dst_pts = np.float32([rkps[m.trainIdx].pt for m in rmatches]).reshape(-1, 1, 2)

            E, mask = cv2.findEssentialMat(dst_pts, src_pts, self.rk, method=cv2.RANSAC, prob=0.999, threshold=1.0)
            src_pts_filtered = src_pts[mask.ravel() == 1]
            dst_pts_filtered = dst_pts[mask.ravel() == 1]
            src_pts_filtered = np.round(src_pts_filtered).astype(np.int16).squeeze()
            dst_pts_filtered = np.round(dst_pts_filtered).astype(np.int16).squeeze()

            _, R, t, mask = cv2.recoverPose(E, dst_pts_filtered, src_pts_filtered, self.lk)

            self.rpose[:3, :3] = np.dot(R, self.rpose[:3, :3])
            self.rpose[:3, 3] += self.rpose[:3, :3].dot(t.squeeze())
            self.rpath.append(self.rpose[:3, 3].copy())

        self.prev_lkps = lkps
        self.prev_rkps = rkps
        self.prev_ldes = ldes
        self.prev_rdes = rdes
        print('pose estimation',time.time()-temps)
        temps = time.time()
        self.compute_xyz_camera()
        print('3d recove',time.time()-temps)
        self.recove()
    def compute_xyz_camera(self):
        # 相机参数（内参矩阵）
        K_left = np.array([[748.3, 0, 490.6,0],
                            [0, 748.3, 506.3,0],
                            [0, 0, 1,0]])

        K_right = np.array([[743.7, 0, 495.3,1],
                           [0, 743.4, 514.7,0],
                           [0, 0, 1,0]])

        # 特征点匹配
        matches = self.bf.match(self.prev_ldes, self.prev_rdes)
        matches = sorted(matches, key=lambda x: x.distance)

        N = 100
        good_matches = matches[:N]
        # 三维重建
        points3D = []
        points2D_left = []
        points2D_right = []
        for match in good_matches:
            idx1 = match.queryIdx
            idx2 = match.trainIdx

            # 左右图像中的特征点
            p1_left = self.prev_lkps[idx1].pt
            p2_right = self.prev_rkps[idx2].pt

            # 对极几何约束
            # F, mask = cv2.findFundamentalMat(np.array([p1_left]), np.array([p2_right]), cv2.FM_LMEDS)
            # print(p1_left,p2_right)
            p1_left = np.hstack((p1_left, 1))
            p2_right = np.hstack((p2_right, 1))

            # 三角化
            point_4d = cv2.triangulatePoints(K_left, K_right, p1_left[:2], p2_right[:2])
            point_3d = point_4d[:3] / point_4d[3]

            points3D.append(point_3d)
            points2D_left.append(p1_left)
            points2D_right.append(p2_right)

        self.points3D = points3D
        self.points2D_left = points2D_left
        self.points2D_right = points2D_right

    def recove(self):
        for pt in self.points3D:
            P_c = np.vstack((pt,1))

            # 构造相机到世界坐标系的变换矩阵 T，假设世界坐标系原点位置为 (0, 0, 0)
            T = np.eye(4)  # 单位矩阵
            T[:3, :3] = self.lpose[:3,:3]
            T[:3, 3] = self.lpose[:3,3]

            # 将相机位姿与相机坐标系下的点相乘，得到世界坐标系下的三维坐标
            P_w = np.dot(T, P_c)
            self.map_pts.append(P_w[:-1])
if __name__ == '__main__':
    steroslam = stero_SLAM()
    win1 = Window_3d()
    read_imgs(win1,steroslam)